"""
Multi-Domain Geometry Module

Port of the Python MultiDomain class to Julia with enhanced functionality.
Handles multi-domain decomposition for spectral element method simulations
with support for 2D and 3D geometries, adaptive refinement, and complex
boundary conditions.
"""

"""
    create_multidomain_impl(n_block::Int, n::Int, time_param::Float64, dim::Int,
                           length::Float64, adaptive_refinement::Bool, 
                           refinement_levels::Int) -> MultiDomain

Create a multi-domain configuration for SEM simulations.
"""
function create_multidomain_impl(n_block::Int, n::Int, time_param::Float64, dim::Int,
                                length::Float64, adaptive_refinement::Bool, 
                                refinement_levels::Int)
    
    # Generate domain distribution
    matrix_dist, g, gi, index_sum = generate_domain_distribution(n_block, dim)
    
    # Create refinement map if requested
    refinement_map = nothing
    corner_points = Tuple{Int,Int}[]
    
    if adaptive_refinement
        # Create L-shaped domain for refinement demonstration
        matrix_dist, corner_points = define_lshape_domain(n_block, dim)
        refinement_map = create_refinement_map(matrix_dist, corner_points, refinement_levels, dim)
        # Recalculate distribution for L-shaped domain
        active_blocks = findall(x -> x > 0, matrix_dist)
        gi = create_index_mapping(matrix_dist)
        
        # Direct check and conversion to avoid Float64 issues
        if isa(active_blocks, Number)
            @warn "findall returned Number instead of array, using empty array"
            active_blocks = []
        elseif isa(active_blocks, AbstractArray)
            # Ensure it's a proper array
            active_blocks = collect(active_blocks)
        else
            @warn "findall returned unexpected type $(typeof(active_blocks)), using empty array"
            active_blocks = []
        end
        
        println("DEBUG: gi type: $(typeof(gi)), value: $gi")
        if isa(gi, Number)
            @warn "create_index_mapping returned Number instead of array, using zeros"
            gi = zeros(Int, size(matrix_dist))
        end
        
        println("DEBUG: About to call Base.length(active_blocks)")
        println("DEBUG: active_blocks type: $(typeof(active_blocks)), value: $active_blocks")
        index_sum = Base.length(active_blocks)
    end
    
    # Use the appropriate g/active_blocks based on whether we did adaptive refinement
    final_g = adaptive_refinement ? active_blocks : g
    
    return MultiDomain{dim}(
        n, n_block, dim, n+1, length, time_param,
        matrix_dist, tuple(final_g...), gi, index_sum,
        adaptive_refinement, refinement_map, corner_points
    )
end

"""
    generate_domain_distribution(n_block::Int, dim::Int, single::Int=1, 
                                ratio::Int=4, seed::Int=1207, nblock::Int=2)

Generate distribution of domains for multi-domain setup.
"""
function generate_domain_distribution(n_block::Int, dim::Int, single::Int=1, 
                                     ratio::Int=4, seed::Int=1207, nblock::Int=2)
    Random.seed!(seed)
    
    # Create random matrix based on dimensionality
    if dim == 3 && single == 1
        random_matrix = rand(0:ratio-1, n_block, n_block, 1)
        matrix_dist = zeros(Int, n_block + 2, n_block + 2, 3)
    else
        random_matrix = rand(0:ratio-1, ntuple(_ -> n_block, dim)...)
        matrix_dist = zeros(Int, ntuple(_ -> n_block + 2, dim)...)
    end
    
    # Apply block modifications (ported from Python logic)
    if nblock == 1
        if dim >= 2
            random_matrix[1, 2] = 0
            random_matrix[2, 2] = 0
        end
    elseif nblock == 2
        if dim >= 2
            random_matrix[2, 2] = 0
        end
    end
    
    # Fill distribution matrix (skip boundaries)
    slices = ntuple(i -> 2:(n_block+1), dim)
    matrix_dist[slices...] = random_matrix
    
    # Remove independent ones (isolated blocks)
    index_sum = count(x -> x > 0, matrix_dist)
    if index_sum > 1
        matrix_dist = remove_independent_ones(matrix_dist, dim)
        index_sum = count(x -> x > 0, matrix_dist)
    end
    
    # Find non-zero elements and create index mapping
    g = findall(x -> x > 0, matrix_dist)
    gi = create_index_mapping(matrix_dist)
    
    return matrix_dist, g, gi, index_sum
end

"""
    remove_independent_ones(matrix::Array{Int}, dim::Int) -> Array{Int}

Remove isolated blocks (no adjacent non-zero elements).
"""
function remove_independent_ones(matrix::Array{Int}, dim::Int)
    result = copy(matrix)
    
    if dim == 2
        for i0 in 2:(size(matrix, 1) - 1)
            for i1 in 2:(size(matrix, 2) - 1)
                if matrix[i0, i1] >= 1
                    neighbors_sum = (matrix[i0-1, i1] + matrix[i0+1, i1] + 
                                   matrix[i0, i1-1] + matrix[i0, i1+1])
                    if neighbors_sum == 0
                        result[i0, i1] = 0
                    else
                        result[i0, i1] = 1
                    end
                end
            end
        end
    elseif dim == 3
        for i0 in 2:(size(matrix, 1) - 1)
            for i1 in 2:(size(matrix, 2) - 1)
                for i2 in 2:(size(matrix, 3) - 1)
                    if matrix[i0, i1, i2] >= 1
                        neighbors_sum = (matrix[i0-1, i1, i2] + matrix[i0+1, i1, i2] + 
                                       matrix[i0, i1-1, i2] + matrix[i0, i1+1, i2] + 
                                       matrix[i0, i1, i2-1] + matrix[i0, i1, i2+1])
                        if neighbors_sum == 0
                            result[i0, i1, i2] = 0
                        else
                            result[i0, i1, i2] = 1
                        end
                    end
                end
            end
        end
    end
    
    return result
end

"""
    create_index_mapping(matrix_dist::AbstractArray{Int}) -> Array{Int}

Create cumulative index mapping for active blocks.
"""
function create_index_mapping(matrix_dist::AbstractArray{Int})
    # Ensure we have a proper integer array
    if !(eltype(matrix_dist) <: Integer)
        matrix_dist = convert(Array{Int}, matrix_dist)
    end
    
    try
        gi = zeros(Int, size(matrix_dist))
        vector = reshape(matrix_dist, :)
        cumsum_vector = cumsum(vector .!= 0) .* (vector .!= 0)
        gi = reshape(cumsum_vector, size(matrix_dist))
        return gi
    catch e
        @warn "create_index_mapping failed: $e, returning zeros"
        return zeros(Int, size(matrix_dist))
    end
end

"""
    calculate_boundary_conditions!(multidomain::MultiDomain{D}, g, gi, 
                                   un_rhs::Array, u2::Array) where D

Calculate boundary conditions for multi-domain setup.
"""
function calculate_boundary_conditions!(multidomain::MultiDomain{D}, g, gi, 
                                       un_rhs::Union{Array,Nothing}, u2::Array) where D
    if un_rhs === nothing
        return un_rhs
    end
    
    un_rhs_orig = copy(un_rhs)
    
    # Process each active block
    for l in 1:multidomain.index_sum
        if D == 3
            i0, i1, i2 = g[l][1], g[l][2], g[l][3]
            process_boundary_3d!(multidomain, gi, un_rhs, un_rhs_orig, u2, l, i0, i1, i2)
        else
            i0, i1 = g[l][1], g[l][2]
            process_boundary_2d!(multidomain, gi, un_rhs, un_rhs_orig, u2, l, i0, i1)
        end
    end
    
    return un_rhs
end

"""
    process_boundary_2d!(multidomain::MultiDomain{2}, gi, un_rhs, un_rhs_orig, 
                         u2, l::Int, i0::Int, i1::Int)

Helper method to process 2D boundary conditions.
"""
function process_boundary_2d!(multidomain::MultiDomain{2}, gi, un_rhs, un_rhs_orig, 
                             u2, l::Int, i0::Int, i1::Int)
    # Get indices of adjacent blocks
    l0 = gi[i0-1, i1] - 1
    l1 = gi[i0, i1-1] - 1
    
    # X- boundary
    if multidomain.matrix_dist[i0-1, i1] == 0
        un_rhs[l+1, 1, :] = u2[l+1, 1, :]
    else
        if l0 >= 0
            un_rhs[l+1, 1, 2:end-1] .+= un_rhs_orig[l0+1, end, 2:end-1]
        end
    end
    
    # Y- boundary
    if multidomain.matrix_dist[i0, i1-1] == 0
        un_rhs[l+1, :, 1] = u2[l+1, :, 1]
    else
        if l1 >= 0
            un_rhs[l+1, 2:end-1, 1] .+= un_rhs_orig[l1+1, 2:end-1, end]
        end
    end
    
    # X+ boundary
    if multidomain.matrix_dist[i0+1, i1] == 0
        un_rhs[l+1, end, :] .= 0
    end
    
    # Y+ boundary
    if multidomain.matrix_dist[i0, i1+1] == 0
        un_rhs[l+1, :, end] .= 0
    end
end

"""
    process_boundary_3d!(multidomain::MultiDomain{3}, gi, un_rhs, un_rhs_orig, 
                         u2, l::Int, i0::Int, i1::Int, i2::Int)

Helper method to process 3D boundary conditions.
"""
function process_boundary_3d!(multidomain::MultiDomain{3}, gi, un_rhs, un_rhs_orig, 
                             u2, l::Int, i0::Int, i1::Int, i2::Int)
    # Get indices of adjacent blocks
    l0 = gi[i0-1, i1, i2] - 1
    l1 = gi[i0, i1-1, i2] - 1
    l2 = gi[i0, i1, i2-1] - 1
    
    # X- boundary
    if multidomain.matrix_dist[i0-1, i1, i2] == 0
        un_rhs[l+1, 1, :, :] = u2[l+1, 1, :, :]
    else
        if l0 >= 0
            un_rhs[l+1, 1, 2:end-1, 2:end-1] .+= un_rhs_orig[l0+1, end, 2:end-1, 2:end-1]
        end
    end
    
    # Y- boundary
    if multidomain.matrix_dist[i0, i1-1, i2] == 0
        un_rhs[l+1, :, 1, :] = u2[l+1, :, 1, :]
    else
        if l1 >= 0
            un_rhs[l+1, 2:end-1, 1, 2:end-1] .+= un_rhs_orig[l1+1, 2:end-1, end, 2:end-1]
        end
    end
    
    # Z- boundary  
    if multidomain.matrix_dist[i0, i1, i2-1] == 0
        un_rhs[l+1, :, :, 1] = u2[l+1, :, :, 1]
    else
        if l2 >= 0
            un_rhs[l+1, 2:end-1, 2:end-1, 1] .+= un_rhs_orig[l2+1, 2:end-1, 2:end-1, end]
        end
    end
    
    # X+, Y+, Z+ boundaries (set to zero for outflow)
    if multidomain.matrix_dist[i0+1, i1, i2] == 0
        un_rhs[l+1, end, :, :] .= 0
    end
    if multidomain.matrix_dist[i0, i1+1, i2] == 0
        un_rhs[l+1, :, end, :] .= 0
    end
    if multidomain.matrix_dist[i0, i1, i2+1] == 0
        un_rhs[l+1, :, :, end] .= 0
    end
end

"""
    construct_multidomain_field(multidomain::MultiDomain{D}, value) where D

Construct a multi-domain array from a scalar or array value.
"""
function construct_multidomain_field(multidomain::MultiDomain{D}, value) where D
    if isa(value, Number)
        return fill(value, multidomain.index_sum)
    end
    
    NP1 = multidomain.n_plus_1
    result_size = (multidomain.index_sum, ntuple(_ -> NP1, D)...)
    result = zeros(eltype(value), result_size)
    
    for l in 1:multidomain.index_sum
        slices = (l, ntuple(_ -> Colon(), D)...)
        result[slices...] = value
    end
    
    return result
end

"""
    construct_multidomain_velocity_fields(multidomain::MultiDomain{D}, 
                                         u, v, w, p, FU, FV, FW) where D

Construct multi-domain arrays for all velocity and pressure fields.
"""
function construct_multidomain_velocity_fields(multidomain::MultiDomain{D}, 
                                              u, v, w, p, FU, FV, FW) where D
    ut = construct_multidomain_field(multidomain, u)
    vt = construct_multidomain_field(multidomain, v)
    wt = construct_multidomain_field(multidomain, w)
    pt = construct_multidomain_field(multidomain, p)
    FUt = construct_multidomain_field(multidomain, FU)
    FVt = construct_multidomain_field(multidomain, FV)
    FWt = construct_multidomain_field(multidomain, FW)
    
    return ut, vt, wt, pt, FUt, FVt, FWt
end

"""
    build_multidomain_grid(multidomain::MultiDomain{D}, xx, yy, zz=nothing) where D

Build coordinate grids for multi-domain setup.
"""
function build_multidomain_grid(multidomain::MultiDomain{D}, xx, yy, zz=nothing) where D
    if D == 2
        return xx, yy, nothing
    else
        return xx, yy, zz
    end
end

"""
    print_multidomain_info(multidomain::MultiDomain{D}) where D

Print information about the multi-domain configuration.
"""
function print_multidomain_info(multidomain::MultiDomain{D}) where D
    println("Multi-Domain Configuration:")
    println("  Dimension: $(D)D")
    println("  Polynomial order: $(multidomain.n)")
    println("  Blocks per dimension: $(multidomain.n_block)")
    println("  Active blocks: $(multidomain.index_sum)")
    println("  Domain length: $(multidomain.length)")
    
    if multidomain.refinement_map !== nothing
        max_level = maximum(multidomain.refinement_map[multidomain.refinement_map .> 0])
        println("  Adaptive refinement: enabled (max level: $max_level)")
        println("  Corner points: $(length(multidomain.corner_points))")
    else
        println("  Adaptive refinement: disabled")
    end
end